# cython: profile=True
from numpy import dot,zeros,array,ones,argmax,zeros_like,shape,isnan
from math import sqrt,sin,cos,exp,pi
import cython
if not cython.compiled:
    from math import *
from v3dsolver.fluid_elements.fluid_elems import Saffman
from v3dsolver.neighbour_finder import Bin,NeighbourFinder
from v3dsolver.vector import Vec
import lpsolver
from random import choice
from mpi4py import MPI
    
class VRTSolver:
    """
    This class is responsible for solving VRT around an element and providing
    vorticities obtained by the neighbours of that element.
    """
    def __init__(self,base_elem,neighbours,len_elems,h_nu,n=8):
        """
        Initialization of VRTSolver class. 
        
        base_elem: element for which VRT is to be solved (of form: [hash,elem])
        neighbours: list of neighbours to base_elem (each of form: [hash,elem])
        len_elems: number of elements in the bin
        h_nu: criteria for neighbourhood: 
             sqrt(8*nu*dt): for 3D 
             sqrt(12*nu*dt): for 2D
        n: number of new possible neighbours on each latitude (same as no. of 
          such latitudes)
        """
        self.base_elem = base_elem
        self.base_index = int(self.base_elem[0].split()[1]) # base_elem index
        self.nbrs = neighbours
        self.h_nu = h_nu
        self.sqrt_nudt = self.h_nu/sqrt(8) ## only for 3D
        self.n = n
        self.possible_nbrs = self.get_possible_nbrs(n) 
        self.len_elems = len_elems
        self._hash_prefix = str(MPI.COMM_WORLD.rank) + ' '
    
    def get_possible_nbrs(self,n):
        """
        Returns a list of uniformly spaced points on a sphere centered at base
        elem and of radius h_nu. 
        
        Any new neighbour would be placed at one of these lat-long positions, 
        depending on distance of this new neighbour from other neighbours.
        
        n: number of new points on each latitude (same as no. of such latitudes)
        """
        pts = []
        for i in range(n/2+1):
            theta=(2.0*pi/n)*i
            z = cos(theta)
            if i==0 or i==(n-1):
                # just add one particle at the pole
                x,y = 0.0,0.0
                pts += [Vec(x,y,z)*self.h_nu+self.base_elem[1].pos]
            else:
                for j in range(n):
                    phi=(2*pi/n)*j
                    x=sin(theta)*cos(phi)
                    y=sin(theta)*sin(phi)
                    pts += [Vec(x,y,z)*self.h_nu+self.base_elem[1].pos]
        return pts
    
    def add_neighbour(self):
        """
        Adds a neighbour particle to base_elem in case the number of neighbours
        it already has is less than the number of equations of VRT. 
        
        It tries to add the neighbour such that the sum of its distance from 
        all the other neighbours is maximum. This approach tends to fill up 
        holes in the vorticity distribution.
        """
        # sum of distance of possible neighbours from all other neighbours
        l = len(self.possible_nbrs)
        #dist_possible_nbrs = zeros(l,dtype='float64')
        
        #for j in range(len(self.nbrs)):
        #    nbr = self.nbrs[j]
        #    nbr_pos = nbr[1]
        #    for i in range(l):
        #        possible_nbr = self.possible_nbrs[i]
        #        rvec = possible_nbr - nbr_pos
        #        r = rvec.length()
        #        dist_possible_nbrs[i] += r 
        #argnbr = argmax(dist_possible_nbrs)
        argnbr = choice(range(l))
       
        new_pos = self.possible_nbrs[argnbr]
        new_index = self.len_elems + self.base_index
        self.base_index = new_index
        str_new_index = str(new_index)
        new_hash = self._hash_prefix + str_new_index
        self.nbrs.append([new_hash, new_pos])

    def get_fractions(self):
        """
        Returns the fractions of vorticity obtained by solving the Vorticity 
        Redistribution Technique equations.
        """
        neq = 10              # number of VRT equations
        nvar = len(self.nbrs) # number of variables = number of neighbours
        Aeq = zeros((neq,nvar))
        id = ones(nvar)
        beq = zeros(neq)
        
        eps1 = zeros(nvar,float)
        eps2 = zeros(nvar,float)
        eps3 = zeros(nvar,float)
        
        for i in range(nvar):
            eps=(self.base_elem[1].pos - self.nbrs[i][1])/self.sqrt_nudt
            eps1[i]=eps[0]
            eps2[i]=eps[1]
            eps3[i]=eps[2]

        # making equation in terms of f_ij
        Aeq[0][:nvar] = id
        Aeq[1][:nvar] = eps1
        Aeq[2][:nvar] = eps2
        Aeq[3][:nvar] = eps3
        Aeq[4][:nvar] = eps1*eps2
        Aeq[5][:nvar] = eps2*eps3
        Aeq[6][:nvar] = eps1*eps3
        Aeq[7][:nvar] = eps1*eps1
        Aeq[8][:nvar] = eps2*eps2
        Aeq[9][:nvar] = eps3*eps3
        
        beq[0] = 1        
        beq[7] = 2
        beq[8] = 2
        beq[9] = 2
        beq1 = zeros(neq)
        for i in range(neq):
            beq1[i] = beq[i] - 0.5*sum(Aeq[i])
        
        w = lpsolver.solver(Aeq,beq1)
        f = w + 0.5*ones(nvar+1) # dimensions of f and w one more than no. of 
                                 # vars (nvar) (last entry stores max value)
        f[-1]=max(f[:nvar])
        return f
    
    def solve_vrt(self):
        """
        Solves the VRT and returns a list of neighbours in the form:
        [nbrhash,nbrpos,nbrvor]; where nbrvor is component of vorticity obtained
        by diffusing current base_elem.
        """
        while len(self.nbrs)<11:
            self.add_neighbour()
        f = self.get_fractions()
        while f[-1]>1 or isnan(f[-1]):
            self.add_neighbour()
            f = self.get_fractions()
        result = []
        assert len(self.nbrs) == len(f) - 1
        for i,nbr in enumerate(self.nbrs):
            result += [[nbr[0],nbr[1],f[i]*self.base_elem[1].vor]]
        return result

class DiffusionManager:
    """
    This class is responsible for diffusion of vorticity of all elements. It 
    uses Vorticity Redistribution Technique (VRT) to simulate diffusion.
    """
    def __init__(self,elems,nu,dt,R=sqrt(8)):
        """
        Initialization of DiffusionManager object.
        
        elems: list of elements to be diffused (only local elems in case of 
               parallel implementation)
        nu: kinematic viscosity of the fluid
        dt: timestep for simulation
        R: radius for binning (sqrt(8) for 3D flow; sqrt(12) for 2D flow)
        """
        self.elems = elems
        self.nu = nu
        self.dt = dt
        self.R = R
        self.h_nu = self.R*sqrt(nu*dt)

    def dispose_nbrs(self,new_nbrs,comm=MPI.COMM_WORLD):
        """
        Dispose obtained neighbours to appropriate locations. 
        
        * Remote processor vorticities are sent back
        * Vorticities obtained from remote processor are updated 
        * Already existing elems update their vorticities according to local 
          diffusion
        * Newly created elems are added to the same processor
        """
        send_data = []
        for i in range(comm.size):
            send_data += [[]]
        
        newelems = []
        newvors = [Vec() for elem in self.elems]
        
        for new_nbr in new_nbrs:
            new_hash,new_pos,new_vor = new_nbr
            
            # if belonging to remote processor
            if int(new_hash.split()[0])!=comm.rank:
                send_data[int(new_hash.split()[0])] += [[new_hash,new_vor]]
            
            # if newly created
            elif int(new_hash.split()[1])>=len(self.elems):
                newelems += [Saffman(pos=new_pos,vel=array([0.,0.,0.]),
                             vor=new_vor,rad=self.elems[0].rad)]
            
            # if already existing
            else:
                index = int(new_hash.split()[1])
                newvors[index] += new_vor
        
        # send non-local data
        reqlist = []
        for i in range(len(send_data)):
            if not i==comm.rank:
                reqlist += [comm.isend(send_data[i],dest=i,tag=7)]
        
        # recv non-local data
        for i in range(comm.size):
            if not i==comm.rank:
                recvd_data = comm.recv(source=i,tag=7)
                for remote_nbr in recvd_data:
                    hash,vor = remote_nbr
                    assert int(hash.split()[0])==comm.rank
                    index = int(hash.split()[1])
                    newvors[index] += vor
        MPI.Prequest.Waitall(reqlist)
        # updating elem vors based on obtained (local and non-local) vors
        for i,elem in enumerate(self.elems):
            elem.vor = newvors[i]
        
        # adding newly created elems to self.elems list
        self.elems += newelems

    def diffuse(self,comm=MPI.COMM_WORLD):
        """
        Perform diffusion using VRT. Returns the list of elems obtained after
        diffusion.
        """
        self.bin = Bin(self.elems,self.h_nu)
        self.nbrfinder = NeighbourFinder(self.bin,self.nu,self.dt,self.R)
        new_nbrs = []
        len_elems = len(self.elems)
        for i,elem in enumerate(self.elems):
            nbrs = self.nbrfinder.get_neighbours_parallel(elem,comm)
            base_elem = [self.bin.hashdict[elem],elem]
            vrtsolver = VRTSolver(base_elem,nbrs,len_elems,self.h_nu)
            vrt_nbrs = vrtsolver.solve_vrt()
            new_nbrs += vrt_nbrs
        self.dispose_nbrs(new_nbrs,comm)
